﻿using System;
using System.IdentityModel.Selectors;
using System.IdentityModel.Tokens;
using System.IO;
using System.Security.Cryptography;
using System.Security.Cryptography.X509Certificates;
using System.ServiceModel.Security;
using System.ServiceModel.Security.Tokens;
using System.Xml;
using System.ServiceModel;

namespace SecurityTokenServiceNS
{
	public class EncryptedWSSecurityTokenSerializer : WSSecurityTokenSerializer
	{
		private X509SecurityToken _rpToken;
		private string _encryptingAlgorithm;// = Constants.XMLEncryption.Algorithm.aes256_cbc;
		private string _asymmetricKeyWrapAlgorithm = Constants.XMLEncryption.Algorithm.rsa_oaep_mgf1p;
		private SecurityAlgorithmSuite _symmetricAlgorithmSuite = SecurityAlgorithmSuite.Basic256;
		private SecurityTokenSerializer _tokenSerializer;

		public EncryptedWSSecurityTokenSerializer(RST rst)
		{
			_rpToken = (X509SecurityToken)StsConfiguration.Current.GetRpToken(rst.AppliesTo);
			_encryptingAlgorithm = rst.EncryptionAlgorithm != "" ? 
				rst.EncryptionAlgorithm : Constants.XMLEncryption.Algorithm.aes256_cbc;
			_tokenSerializer = new WSSecurityTokenSerializer();
		}

		protected override void WriteTokenCore(XmlWriter writer, SecurityToken token)
		{
			XmlElement encryptedToken;

			using (MemoryStream ms = new MemoryStream())
			{
				XmlDictionaryWriter wr = XmlDictionaryWriter.CreateDictionaryWriter(new XmlTextWriter(new StreamWriter(ms)));
				_tokenSerializer.WriteToken(wr, token);
				wr.Flush();
				ms.Seek(0L, SeekOrigin.Begin);
				encryptedToken = Encrypt(ms);
				wr.Close();
				ms.Close();
			}

			encryptedToken.WriteTo(writer);
		}

		private XmlElement Encrypt(MemoryStream encryptStream)
		{
			SecurityKeyIdentifier rpTokenReference =
				new SecurityKeyIdentifier(new SecurityKeyIdentifierClause[]
					{ _rpToken.CreateKeyIdentifierClause<X509ThumbprintKeyIdentifierClause>() }	);

			RNGCryptoServiceProvider random = new RNGCryptoServiceProvider();
			byte[] buffer = new byte[_symmetricAlgorithmSuite.DefaultEncryptionKeyDerivationLength / 8];
			random.GetNonZeroBytes(buffer);

			WrappedKeySecurityToken token = new WrappedKeySecurityToken(string.Empty, buffer, _asymmetricKeyWrapAlgorithm, _rpToken, rpTokenReference);
			SecurityKeyIdentifier identifier = new SecurityKeyIdentifier(new SecurityKeyIdentifierClause[]
				{ new EncryptedKeyIdentifierClause(token.GetWrappedKey(), token.WrappingAlgorithm, token.WrappingTokenReference) });
			SymmetricAlgorithm symmetricAlgorithm = ((SymmetricSecurityKey)token.SecurityKeys[0]).GetSymmetricAlgorithm(_encryptingAlgorithm);
			EncryptedData data = new EncryptedData(identifier, _encryptingAlgorithm, symmetricAlgorithm, encryptStream.GetBuffer(), encryptStream.Length);

			MemoryStream stream = new MemoryStream();
			XmlWriter writer = XmlWriter.Create(stream);
			data.WriteTo(writer);
			writer.Flush();

			XmlDocument doc = new XmlDocument();
			doc.PreserveWhitespace = false;
			stream.Position = 0;
			doc.Load(stream);

			stream.Close();
			return doc.DocumentElement;
		}
	}
}
